#include <iostream>
#include <chrono>
#include <omp.h>
#include <arm_neon.h>

struct ListNode{
    int ptr_id_;
    int cache_line_offset_;
    ListNode *next;
    ListNode(){};
    ListNode(int ptr_id, int cache_line_offset): ptr_id_(ptr_id), cache_line_offset_(cache_line_offset), next(nullptr){};
};

struct EvalCacheHook{
    ListNode *cache_;
    std::unordered_map<int, size_t> perfRes_;
    size_t recordSz_;
    int cacheSz_;
    int cacheLineSz_;
    EvalCacheHook(ListNode *cache, int m, int n, int k, int cacheSz): cache_(cache), perfRes_(), recordSz_(0), cacheSz_(cacheSz), cacheLineSz_(8){}
    void free(){
        dumpPerf();
        while (cache_)
        {
            auto leave = cache_->next;
            delete cache_;
            cache_ = leave;
        }
    }

    void dumpPerf(){
        if(perfRes_.empty()) return;
        std::cout << " perf_start " << std::endl;
        int reuseDistance;
        size_t times;
        float reuseDistanceSum = 0.f;
        int sum_query = 0;
        int miss_cache = 0;
        for(auto &pair : perfRes_){
            std::tie(reuseDistance, times) = pair;
            std::cout << " " << reuseDistance << " " << times << std::endl;
            if(reuseDistance == -1){
                reuseDistance = cacheSz_*100;
                miss_cache = times;
            }
            sum_query += times;
            reuseDistanceSum += float(reuseDistance)*float(times);
        }
        std::cout << " sum_query " << sum_query << " miss_cache " << miss_cache << " miss_ratio " << float(miss_cache)/float(sum_query)
            << " reuse_distance_sum " << reuseDistanceSum << std::endl;
    }
    void samplePerf(int reuseDistance){
        if(reuseDistance > cacheSz_){
            reuseDistance = -1;
        }
        if(perfRes_.find(reuseDistance) == perfRes_.end()){
            perfRes_[reuseDistance] = 0;
        }
        ++perfRes_[reuseDistance];
    }

    void record(int ptr_id, int ptr_offset){
        // ++recordSz_;
        // if(recordSz_ % 100000 == 0){
        //     std::cout << " record " << ptr_id << " " << ptr_offset << std::endl;
        // }
        int cache_line_offset = ptr_offset/cacheLineSz_;
        auto cur = cache_->next, pre = cache_;
        ListNode *target = nullptr;
        int reuseDistance = 0;
        while (cur)
        {
            if(cur->ptr_id_ == ptr_id && cur->cache_line_offset_ == cache_line_offset){
                target = cur;
                pre->next = cur->next;
                target->next = nullptr;
                target->next = cache_->next;
                cache_->next = target;
                samplePerf(reuseDistance);
                break;
            }
            pre = cur;
            cur = cur->next;
            ++reuseDistance;
        }
        if(!target){
            samplePerf(-1);
            auto tmp = new ListNode(ptr_id, cache_line_offset);
            tmp->next = cache_->next;
            cache_->next = tmp;
        }
    }
};


static std::uint64_t now() {
  return std::chrono::duration_cast<std::chrono::milliseconds>
    (std::chrono::high_resolution_clock::now().time_since_epoch()).count();
}

template<typename T>
void set(T *src, int sz, const T val){
    for(int i = 0; i < sz; ++i){
        src[i] = val;
    }
}

template<typename T>
void eval(T *src, T *gt, int sz){
    for(int i = 0; i < sz; ++i){
        if(src[i] != gt[i]){
            std::cout << "!!! no equal pos " << i << " src " << src[i] << " gt " << gt[i] << std::endl; 
            break;
        }else{
            // if(i % 1000 == 0) std::cout << " equal " << src[i] << std::endl;
        }
    }
}

template<typename T>
void gemm_gt(T *A, T *B, T *C, int M, int N, int K, EvalCacheHook *evalCache = nullptr){
    for(int m = 0; m < M; ++m){
        for(int k = 0; k < K; ++k){
            for(int n = 0; n < N; ++n){
                #ifdef DEBUG
                if(evalCache){
                    evalCache->record(0, m*K+k);
                    evalCache->record(1, k*N+n);
                    evalCache->record(2, m*N+n);
                }
                #endif
                #ifdef TEST
                C[m*N+n] += A[m*K+k]*B[k*N+n];
                #endif
            }
        }
    }
}

template<typename T>
void transpose(T *A, T *B, const int m, const int n){
    for(int m_id = 0; m_id < m; ++m_id){
        for(int n_id = 0; n_id < n; ++n_id){
            B[n_id*m+m_id] = A[m_id*n+n_id];
        }
    }
}

//row major
inline void fmma_4x4(float *A, float *B, float *C, const int m, const int n, const int k){ 
    // these are the rows A
    float32x4_t A0;
    float32x4_t A1;
    float32x4_t A2;
    float32x4_t A3;
    
    // these are the rows B
    float32x4_t B0;
    float32x4_t B1;
    float32x4_t B2;
    float32x4_t B3;
    
    // these are the rows C
    float32x4_t C0;
    float32x4_t C1;
    float32x4_t C2;
    float32x4_t C3;

    B0 = vld1q_f32(B);
    B1 = vld1q_f32(B+n);
    B2 = vld1q_f32(B+2*n);
    B3 = vld1q_f32(B+3*n);
    
    C0 = vld1q_f32(C);
    C1 = vld1q_f32(C+n);
    C2 = vld1q_f32(C+2*n);
    C3 = vld1q_f32(C+3*n);

    //k = [0, 4), m = 0, n = [0, 4)
    A0 = vld1q_f32(A);
    C0 = vfmaq_laneq_f32(C0, B0, A0, 0);
    C0 = vfmaq_laneq_f32(C0, B1, A0, 1);
    C0 = vfmaq_laneq_f32(C0, B2, A0, 2);
    C0 = vfmaq_laneq_f32(C0, B3, A0, 3);
    vst1q_f32(C, C0);

    //k = [0, 4), m = 1, n = [0, 4)
    A1 = vld1q_f32(A+k);
    C1 = vfmaq_laneq_f32(C1, B0, A1, 0);
    C1 = vfmaq_laneq_f32(C1, B1, A1, 1);
    C1 = vfmaq_laneq_f32(C1, B2, A1, 2);
    C1 = vfmaq_laneq_f32(C1, B3, A1, 3);
    vst1q_f32(C+n, C1);

    A2 = vld1q_f32(A+2*k);
    C2 = vfmaq_laneq_f32(C2, B0, A2, 0);
    C2 = vfmaq_laneq_f32(C2, B1, A2, 1);
    C2 = vfmaq_laneq_f32(C2, B2, A2, 2);
    C2 = vfmaq_laneq_f32(C2, B3, A2, 3);
    vst1q_f32(C+2*n, C2);

    A3 = vld1q_f32(A+3*k);
    C3 = vfmaq_laneq_f32(C3, B0, A3, 0);
    C3 = vfmaq_laneq_f32(C3, B1, A3, 1);
    C3 = vfmaq_laneq_f32(C3, B2, A3, 2);
    C3 = vfmaq_laneq_f32(C3, B3, A3, 3);
    vst1q_f32(C+3*n, C3);
}

template<int bm, int bn, int bk, typename T>
void gemm_0(T *A, T *B, T *C, int m, int n, int k, EvalCacheHook *evalCache = nullptr){
    T Ac[bm*bk] = {0};// for cache
    T Bc[bk*bn] = {0};
    for(int bm_start = 0; bm_start < m; bm_start+=bm){
        for(int bn_start = 0; bn_start < n; bn_start+=bn){
            T Cc[bm*bn] = {0};
            for(int bk_start = 0; bk_start < k; bk_start += bk){
                //load block A
                for(int tm_id = 0; tm_id < bm; ++tm_id){
                    int m_id = bm_start+tm_id;
                    for(int tk_id=0; tk_id < bk; ++tk_id){
                        int k_id = bk_start+tk_id;
                        if(m_id < m && k_id < k){
                            Ac[tm_id*bk+tk_id] = A[m_id*k+k_id];
                        }else{
                            Ac[tm_id*bk+tk_id] = 0;
                        }
                    }
                }                
                //load block B
                for(int tk_id=0; tk_id < bk; ++tk_id){
                    int k_id = bk_start+tk_id;
                    for(int tn_id = 0; tn_id < bn; ++tn_id){
                        int n_id = bn_start+tn_id;
                        if(k_id < k && n_id < n){
                            Bc[tk_id*bn+tn_id] = B[k_id*n+n_id];
                        }else{
                            Bc[tk_id*bn+tn_id] = 0;
                        }
                    }
                }
                //math
                for(int tm_id = 0; tm_id < bm; tm_id+=4){
                    for(int tn_id = 0; tn_id < bn; tn_id+=4){
                        // these are the rows C
                        float32x4_t C0, C1, C2, C3;
                        float *Ct = Cc+tm_id*bn+tn_id;
                        C0 = vld1q_f32(Ct);
                        C1 = vld1q_f32(Ct+bn);
                        C2 = vld1q_f32(Ct+2*bn);
                        C3 = vld1q_f32(Ct+3*bn);
                        for(int tk_id=0; tk_id < bk; tk_id+=4){
                            // these are the rows A, rows B
                            float32x4_t A0, A1, A2, A3;
                            float32x4_t B0, B1, B2, B3;
                            float *Bt = Bc+tk_id*bn+tn_id;
                            float *At = Ac+tm_id*bk+tk_id;
                            B0 = vld1q_f32(Bt);
                            B1 = vld1q_f32(Bt+bn);
                            B2 = vld1q_f32(Bt+2*bn);
                            B3 = vld1q_f32(Bt+3*bn);
                        
                            //k = [0, 4), m = 0, n = [0, 4)
                            A0 = vld1q_f32(At);
                            C0 = vfmaq_laneq_f32(C0, B0, A0, 0);
                            C0 = vfmaq_laneq_f32(C0, B1, A0, 1);
                            C0 = vfmaq_laneq_f32(C0, B2, A0, 2);
                            C0 = vfmaq_laneq_f32(C0, B3, A0, 3);
                        
                            //k = [0, 4), m = 1, n = [0, 4)
                            A1 = vld1q_f32(At+bk);
                            C1 = vfmaq_laneq_f32(C1, B0, A1, 0);
                            C1 = vfmaq_laneq_f32(C1, B1, A1, 1);
                            C1 = vfmaq_laneq_f32(C1, B2, A1, 2);
                            C1 = vfmaq_laneq_f32(C1, B3, A1, 3);
                        
                            A2 = vld1q_f32(At+2*bk);
                            C2 = vfmaq_laneq_f32(C2, B0, A2, 0);
                            C2 = vfmaq_laneq_f32(C2, B1, A2, 1);
                            C2 = vfmaq_laneq_f32(C2, B2, A2, 2);
                            C2 = vfmaq_laneq_f32(C2, B3, A2, 3);
                        
                            A3 = vld1q_f32(At+3*bk);
                            C3 = vfmaq_laneq_f32(C3, B0, A3, 0);
                            C3 = vfmaq_laneq_f32(C3, B1, A3, 1);
                            C3 = vfmaq_laneq_f32(C3, B2, A3, 2);
                            C3 = vfmaq_laneq_f32(C3, B3, A3, 3);
                        }
                        vst1q_f32(Ct, C0);
                        vst1q_f32(Ct+bn, C1);
                        vst1q_f32(Ct+2*bn, C2);
                        vst1q_f32(Ct+3*bn, C3);
                    }
                }
            }
            //store C
            for(int tm_id = 0; tm_id < bm; ++tm_id){
                int m_id = bm_start+tm_id;
                for(int tn_id = 0; tn_id < bn; ++tn_id){
                    int n_id = bn_start+tn_id;
                    C[m_id*n+n_id] += Cc[tm_id*bn+tn_id];
                }
            }
        }
    }
}

//FLOPS (3.5*1024*1024*1024)

int main(){
    // const int m = 1000, n = 1000, k = 1000;
    const int m = 512, n = 798, k = 512;
    float *A = new float[m*k];
    float *B = new float[k*n];
    float *C = new float[m*n];
    float *C_gt = new float[m*n];
    float latency_gt = 1.0f, latency_test = 0.f;
    //gt
    {
        set(A, m*k, 1.0f);
        set(B, k*n, 1.0f);
        set(C_gt, m*n, 1.0f);
        auto start_t = now();

        auto evalCache = new EvalCacheHook(new ListNode(-1, -1), m, n, k, 128*256);
        gemm_gt(A, B, C_gt, m, n, k, evalCache);
        evalCache->free();
        delete evalCache;

        auto end_t = now();
        latency_gt = float(end_t-start_t);
    }

    //test
    {
        const int blockSz = 64;
        const int bm = blockSz, bn = blockSz, bk = blockSz/8;
        set(A, m*k, 1.0f);
        set(B, k*n, 1.0f);
        set(C, m*n, 1.0f);

        auto start_t = now();

        auto evalCache = new EvalCacheHook(new ListNode(-1, -1), m, n, k, 128*256);
        gemm_0<bm, bn, bk>(A, B, C, m, n, k, evalCache);
        evalCache->free();
        delete evalCache;

        auto end_t = now();

        latency_test = float(end_t-start_t);
    }
    printf(" m %d n %d k %d latency gt %0.2f test %0.2f ratio %0.2f \n", m, n, k, latency_gt, latency_test, latency_test/latency_gt);

    //eval
    {
        eval(C, C_gt, m*n);
    }
    return 0;
}